# 从typing模块导入必要的类型定义
from typing import TypedDict, List

from langchain.prompts import PromptTemplate
from langchain.schema import HumanMessage
from langchain_openai import ChatOpenAI
from langgraph.graph import StateGraph, END


# 定义一个名为 'State' 的TypedDict，用于结构化表示状态数据
class State(TypedDict):
    text: str  # 存储原始输入文本
    classification: str  # 分类结果（例如类别标签）
    entities: List[str]  # 存储提取的实体列表（如命名实体）
    summary: str  # 存储文本的摘要版本


llm = ChatOpenAI(model="qwen-plus", temperature=0, api_key="",
                 base_url="https://dashscope.aliyuncs.com/compatible-mode/v1")


def classification_node(state: State):
    """
    将文本分类到预定义的类别中。

    参数:
        state (State): 包含待分类文本的当前状态字典

    返回:
        dict: 包含"分类"键的字典，值为分类结果

    分类类别:
        - 新闻: 事实性报道当前事件
        - 博客: 个人或非正式的网络写作
        - 研究: 学术或科学内容
        - 其他: 不符合上述类别的内容
    """

    # 创建提示模板，要求模型将文本分类到指定类别
    prompt = PromptTemplate(
        input_variables=["text"],
        template="将以下文本分类到以下类别之一：新闻、博客、研究、其他。\n\n文本：{text}\n\n类别："
    )

    # 使用状态中的文本格式化提示并包装为HumanMessage
    message = HumanMessage(content=prompt.format(text=state["text"]))

    # 调用语言模型进行分类
    classification = llm.invoke([message]).content.strip()

    # 返回分类结果字典
    return {"classification": classification}


def entity_extraction_node(state: State):
    # 从文本中识别并提取命名实体（按人物、组织、地点分类）

    # 创建实体提取提示模板，要求返回逗号分隔的实体列表
    prompt = PromptTemplate(
        input_variables=["text"],
        template="从以下文本中提取所有实体（人物、组织、地点）。以逗号分隔列表形式返回结果。\n\n文本：{text}\n\n实体："
    )

    # 使用文本格式化提示并包装为HumanMessage
    message = HumanMessage(content=prompt.format(text=state["text"]))

    # 调用模型获取响应，清理空白并拆分为列表
    entities = llm.invoke([message]).content.strip().split(", ")

    # 返回包含实体列表的字典
    return {"entities": entities}


def summarize_node(state: State):
    # 创建摘要提示模板，要求用一句话总结输入文本

    # 使用模板创建链式处理流程（提示模板 → 语言模型）
    summarization_prompt = PromptTemplate.from_template(
        """用一句话总结以下文本。\n\n文本：{text}\n\n摘要："""
    )

    # 构建处理链：提示模板 → 语言模型
    chain = summarization_prompt | llm

    # 执行处理流程，传入文本进行摘要生成
    response = chain.invoke({"text": state["text"]})

    # 返回包含摘要结果的字典
    return {"summary": response.content}


# 创建状态图对象
workflow = StateGraph(State)

# 添加处理节点到流程图
workflow.add_node("classification_node", classification_node)  # 分类节点
workflow.add_node("entity_extraction", entity_extraction_node)  # 实体提取节点
workflow.add_node("summarization", summarize_node)  # 摘要生成节点

# 设置流程入口点并定义执行路径
workflow.set_entry_point("classification_node")  # 设置初始执行节点
workflow.add_edge("classification_node", "entity_extraction")  # 分类 → 实体提取
workflow.add_edge("entity_extraction", "summarization")  # 实体提取 → 摘要生成
workflow.add_edge("summarization", END)  # 摘要生成 → 流程结束

# 编译工作流为可执行应用
app = workflow.compile()

# 测试用例：关于Anthropic的MCP技术说明文本
sample_text = """
Anthropic的MCP（Model Context Protocol）是一个开源的强大工具，允许应用程序与各种系统的API进行无缝交互。
"""

# 创建初始状态对象
state_input = {"text": sample_text}

# 执行完整工作流处理
result = app.invoke(state_input)

# 输出处理结果
print("分类结果:", result["classification"])
print("\n实体列表:", result["entities"])
print("\n摘要内容:", result["summary"])
